{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/embeddings/sagemaker_embedding_endpoint.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interacting with Embeddings deployed in Amazon SageMaker Endpoint with LlamaIndex\n",
    "\n",
    "An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models, for making predictions on new data.\n",
    "\n",
    "This notebook demonstrates how to interact with Embedding endpoints using `SageMakerEmbedding`, unlocking additional llamaIndex features.\n",
    "So, It is assumed that an Embedding is deployed on a SageMaker endpoint."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting Up\n",
    "If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-sagemaker-endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install llama-index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You have to specify the endpoint name to interact with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ENDPOINT_NAME = \"<-YOUR-ENDPOINT-NAME->\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Credentials should be provided to connect to the endpoint. You can either:\n",
    "-  use an AWS profile by specifying the `profile_name` parameter, if not specified, the default credential profile will be used. \n",
    "-  Pass credentials as parameters (`aws_access_key_id`, `aws_secret_access_key`, `aws_session_token`, `region_name`). \n",
    "\n",
    "for more details check [this link](https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**AWS profile name**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding\n",
    "\n",
    "AWS_ACCESS_KEY_ID = \"<-YOUR-AWS-ACCESS-KEY-ID->\"\n",
    "AWS_SECRET_ACCESS_KEY = \"<-YOUR-AWS-SECRET-ACCESS-KEY->\"\n",
    "AWS_SESSION_TOKEN = \"<-YOUR-AWS-SESSION-TOKEN->\"\n",
    "REGION_NAME = \"<-YOUR-ENDPOINT-REGION-NAME->\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_model = SageMakerEmbedding(\n",
    "    endpoint_name=ENDPOINT_NAME,\n",
    "    aws_access_key_id=AWS_ACCESS_KEY_ID,\n",
    "    aws_secret_access_key=AWS_SECRET_ACCESS_KEY,\n",
    "    aws_session_token=AWS_SESSION_TOKEN,\n",
    "    aws_region_name=REGION_NAME,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**With credentials**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.embeddings.sagemaker_endpoint import SageMakerEmbedding\n",
    "\n",
    "ENDPOINT_NAME = \"<-YOUR-ENDPOINT-NAME->\"\n",
    "PROFILE_NAME = \"<-YOUR-PROFILE-NAME->\"\n",
    "embed_model = SageMakerEmbedding(\n",
    "    endpoint_name=ENDPOINT_NAME, profile_name=PROFILE_NAME\n",
    ")  # Omit the profile name to use the default profile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call `get_text_embedding` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = embed_model.get_text_embedding(\n",
    "    \"An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models, specifically LLM (Large Language Models), for making predictions on new data.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.021565623581409454,\n",
       " 0.019147753715515137,\n",
       " -0.056700449436903,\n",
       " -0.022435730323195457,\n",
       " -0.03930153697729111,\n",
       " 0.0028269595932215452,\n",
       " 0.13571178913116455,\n",
       " -0.01545825693756342,\n",
       " 0.008395835757255554,\n",
       " -0.028859995305538177,\n",
       " -0.002176997484639287,\n",
       " -0.01331179216504097,\n",
       " 0.04315004497766495,\n",
       " -0.06485316157341003,\n",
       " 0.05034399777650833,\n",
       " 0.03711436316370964,\n",
       " 0.08230344206094742,\n",
       " -0.025957593694329262,\n",
       " -0.022912224754691124,\n",
       " -0.0353832021355629,\n",
       " 0.057388801127672195,\n",
       " 0.06676290929317474,\n",
       " -0.11050011962652206,\n",
       " -0.06368492543697357,\n",
       " 0.06308538466691971,\n",
       " -0.018852291628718376,\n",
       " 0.03129011392593384,\n",
       " -0.0760543942451477,\n",
       " -0.03160060569643974,\n",
       " 0.018145989626646042,\n",
       " 0.023656774312257767,\n",
       " -0.03639936074614525,\n",
       " 0.017052484676241875,\n",
       " -0.02450205385684967,\n",
       " 0.027692550793290138,\n",
       " 0.0025261295959353447,\n",
       " 0.027999257668852806,\n",
       " -0.01860380358994007,\n",
       " 0.037128496915102005,\n",
       " -0.08015701919794083,\n",
       " 0.06850703060626984,\n",
       " -0.008332699537277222,\n",
       " 0.06535708159208298,\n",
       " 0.0328301340341568,\n",
       " 0.006906852591782808,\n",
       " 0.02144588902592659,\n",
       " 0.016819249838590622,\n",
       " 0.065149687230587,\n",
       " -0.044132206588983536,\n",
       " -0.06759986281394958,\n",
       " -0.06730470061302185,\n",
       " 0.04957125708460808,\n",
       " 0.06852839142084122,\n",
       " 0.025557655841112137,\n",
       " 0.02020672708749771,\n",
       " 0.010702744126319885,\n",
       " -0.08458871394395828,\n",
       " 0.015973811969161034,\n",
       " -0.04692167788743973,\n",
       " 0.006929733324795961,\n",
       " -0.07795898616313934,\n",
       " -0.0068844580091536045,\n",
       " 0.00868065282702446,\n",
       " 0.05663231015205383,\n",
       " 0.0727660208940506,\n",
       " 0.025757266208529472,\n",
       " -0.005326052196323872,\n",
       " -0.01100397203117609,\n",
       " -0.12615707516670227,\n",
       " -0.05008727312088013,\n",
       " 0.011879516765475273,\n",
       " 0.05348730832338333,\n",
       " 0.012098579667508602,\n",
       " -0.09577080607414246,\n",
       " 0.10378099977970123,\n",
       " 0.032600581645965576,\n",
       " 0.030491407960653305,\n",
       " -0.08457385003566742,\n",
       " 0.008462479338049889,\n",
       " 0.0191865935921669,\n",
       " -0.07925017923116684,\n",
       " -0.014723804779350758,\n",
       " -0.030580421909689903,\n",
       " -0.042236607521772385,\n",
       " -0.05877465009689331,\n",
       " 0.010675436817109585,\n",
       " 0.032128408551216125,\n",
       " -0.0562652088701725,\n",
       " 0.02557305432856083,\n",
       " -0.04288778454065323,\n",
       " 0.06932021677494049,\n",
       " -0.0056538935750722885,\n",
       " -0.034259598702192307,\n",
       " -0.03814560920000076,\n",
       " -0.050656333565711975,\n",
       " -0.07511229068040848,\n",
       " -0.0727902203798294,\n",
       " 0.057129207998514175,\n",
       " 0.0038894452154636383,\n",
       " -0.026942580938339233,\n",
       " 0.0033212697599083185,\n",
       " -0.0017185898032039404,\n",
       " 0.09014856070280075,\n",
       " -0.050106242299079895,\n",
       " -0.03977276012301445,\n",
       " 0.09139455854892731,\n",
       " -0.047016214579343796,\n",
       " -0.08531513810157776,\n",
       " 0.028743667528033257,\n",
       " -0.09533917158842087,\n",
       " -0.0395311564207077,\n",
       " 0.05482438951730728,\n",
       " 0.08327925205230713,\n",
       " 0.0976642593741417,\n",
       " 0.009845257736742496,\n",
       " -0.04681103676557541,\n",
       " -0.005461001303046942,\n",
       " -0.008060128428041935,\n",
       " 0.03265204653143883,\n",
       " -0.030966755002737045,\n",
       " 0.08599787205457687,\n",
       " -0.10304504632949829,\n",
       " -0.026105130091309547,\n",
       " -0.059405822306871414,\n",
       " 0.0397467203438282,\n",
       " -0.10144900530576706,\n",
       " -0.08908584713935852,\n",
       " -0.002865528455004096,\n",
       " 0.018968375399708748,\n",
       " -0.015612486749887466,\n",
       " -0.02899557165801525,\n",
       " 0.014865572564303875,\n",
       " 0.02968095801770687,\n",
       " 0.02684851363301277,\n",
       " 0.05083318054676056,\n",
       " 0.02924215793609619,\n",
       " 0.01617436856031418,\n",
       " -0.12093227356672287,\n",
       " -0.03196290135383606,\n",
       " -0.06529344618320465,\n",
       " -0.026738135144114494,\n",
       " -0.006192725617438555,\n",
       " -0.0158796738833189,\n",
       " 0.03231367841362953,\n",
       " 0.0864284336566925,\n",
       " 0.009570070542395115,\n",
       " 0.0716225802898407,\n",
       " -0.07572516798973083,\n",
       " 0.05317925661802292,\n",
       " 0.06873224675655365,\n",
       " -0.0053744143806397915,\n",
       " 0.07709360867738724,\n",
       " 0.04732672870159149,\n",
       " 0.030548326671123505,\n",
       " 0.0027018985711038113,\n",
       " 0.008613159880042076,\n",
       " -0.027862899005413055,\n",
       " -0.047938231378793716,\n",
       " -0.0024136374704539776,\n",
       " 0.10075152665376663,\n",
       " -0.05304369330406189,\n",
       " -0.002246642718091607,\n",
       " -0.029105287045240402,\n",
       " -0.06296665966510773,\n",
       " -0.014757883735001087,\n",
       " -0.052865151315927505,\n",
       " 0.006513366475701332,\n",
       " 0.06221649795770645,\n",
       " -0.011616945266723633,\n",
       " -0.009876491501927376,\n",
       " -0.09408566355705261,\n",
       " 0.031154463067650795,\n",
       " -0.08682437986135483,\n",
       " -0.06839584559202194,\n",
       " -0.004306801129132509,\n",
       " -0.04914987087249756,\n",
       " -0.06964779645204544,\n",
       " -0.07772491127252579,\n",
       " 0.020129483193159103,\n",
       " 0.046311721205711365,\n",
       " -0.05257619544863701,\n",
       " -0.019545624032616615,\n",
       " -0.08561691641807556,\n",
       " -0.0333973690867424,\n",
       " -0.07877060025930405,\n",
       " 0.02237010933458805,\n",
       " 0.04267224296927452,\n",
       " -0.02596845105290413,\n",
       " -0.07537004351615906,\n",
       " -0.09591358155012131,\n",
       " -0.02983330562710762,\n",
       " -0.013941322453320026,\n",
       " -0.011087455786764622,\n",
       " -0.012412399053573608,\n",
       " -0.02150963619351387,\n",
       " 0.07786650210618973,\n",
       " 0.04773382097482681,\n",
       " 0.019922567531466484,\n",
       " 0.07738719880580902,\n",
       " 0.05013704299926758,\n",
       " -0.09917087852954865,\n",
       " 0.06646912544965744,\n",
       " -0.029166480526328087,\n",
       " -0.01662266068160534,\n",
       " 0.052809495478868484,\n",
       " 0.017951449379324913,\n",
       " -0.08008625358343124,\n",
       " 0.07121209800243378,\n",
       " -0.02369554713368416,\n",
       " -0.05586501955986023,\n",
       " 0.0498092845082283,\n",
       " -0.01365865021944046,\n",
       " -0.02708507515490055,\n",
       " 0.08256032317876816,\n",
       " 0.025594351813197136,\n",
       " -0.061894018203020096,\n",
       " 0.012498111464083195,\n",
       " 0.11298499256372452,\n",
       " -0.002179534872993827,\n",
       " 0.04463203623890877,\n",
       " 0.010685713030397892,\n",
       " 0.03099050745368004,\n",
       " 0.05935501679778099,\n",
       " 0.058455295860767365,\n",
       " -0.02524683251976967,\n",
       " 0.07769811153411865,\n",
       " 0.03721567243337631,\n",
       " -0.05615007132291794,\n",
       " 0.01681063137948513,\n",
       " 0.04380310699343681,\n",
       " -0.01826358400285244,\n",
       " -0.008245888166129589,\n",
       " -0.05828167498111725,\n",
       " -0.0558970607817173,\n",
       " 0.0023540062829852104,\n",
       " -0.02012726105749607,\n",
       " -0.0008245687349699438,\n",
       " 0.00417731748893857,\n",
       " 0.10165227949619293,\n",
       " -0.04578690975904465,\n",
       " -0.026630911976099014,\n",
       " -0.0011021343525499105,\n",
       " 0.022034477442502975,\n",
       " 0.028991390019655228,\n",
       " 0.05962913855910301,\n",
       " -0.0630992129445076,\n",
       " -0.06605151295661926,\n",
       " -0.04041984677314758,\n",
       " 0.024546820670366287,\n",
       " 0.09267142415046692,\n",
       " 0.08537988364696503,\n",
       " -0.004139745142310858,\n",
       " 0.005737761966884136,\n",
       " 0.009916293434798717,\n",
       " 0.06632582098245621,\n",
       " 0.06800509989261627,\n",
       " 0.055296897888183594,\n",
       " 0.02332977205514908,\n",
       " -0.0944223552942276,\n",
       " 0.017701460048556328,\n",
       " -0.08671680837869644,\n",
       " -0.03739886358380318,\n",
       " -0.05368325486779213,\n",
       " 0.024832075461745262,\n",
       " -0.0017475909553468227,\n",
       " -0.01282496377825737,\n",
       " -0.02055644989013672,\n",
       " -0.03824482858181,\n",
       " -0.0029880295041948557,\n",
       " -0.06276058405637741,\n",
       " -0.0031180321238934994,\n",
       " 0.018192732706665993,\n",
       " -0.03707187995314598,\n",
       " 0.021848147734999657,\n",
       " 0.02335541695356369,\n",
       " -0.01648322306573391,\n",
       " 0.1435147374868393,\n",
       " 0.09290475398302078,\n",
       " 0.006393723655492067,\n",
       " -0.04305407404899597,\n",
       " 0.07931894063949585,\n",
       " 0.027527527883648872,\n",
       " 0.014633690007030964,\n",
       " -0.0010841790353879333,\n",
       " -0.08786225318908691,\n",
       " -0.058670710772275925,\n",
       " -0.01426162850111723,\n",
       " 0.0070671215653419495,\n",
       " -0.006778248585760593,\n",
       " -0.11174825578927994,\n",
       " -0.0030391064938157797,\n",
       " -0.04667389392852783,\n",
       " 0.003977365326136351,\n",
       " 0.044250018894672394,\n",
       " 0.07857362180948257,\n",
       " 0.02811703085899353,\n",
       " -0.03371573984622955,\n",
       " 0.04170072823762894,\n",
       " 0.04904920980334282,\n",
       " 0.043538808822631836,\n",
       " -0.03201114386320114,\n",
       " -0.10568009316921234,\n",
       " 0.054847411811351776,\n",
       " 0.04661855846643448,\n",
       " -0.014991141855716705,\n",
       " 0.05715333670377731,\n",
       " 0.09562608599662781,\n",
       " 0.07444649189710617,\n",
       " 0.05350692570209503,\n",
       " -0.022917866706848145,\n",
       " 0.05435822904109955,\n",
       " -0.05491911619901657,\n",
       " -0.02363002300262451,\n",
       " -0.006824106443673372,\n",
       " 0.03266216069459915,\n",
       " 0.11281733214855194,\n",
       " -0.03085748292505741,\n",
       " 0.08682505786418915,\n",
       " -0.024763137102127075,\n",
       " -0.04231179505586624,\n",
       " 0.0027543380856513977,\n",
       " -0.05821159482002258,\n",
       " 0.011324281804263592,\n",
       " 0.004590475000441074,\n",
       " 0.024966349825263023,\n",
       " -0.01651289500296116,\n",
       " -0.0320640504360199,\n",
       " 0.02205571159720421,\n",
       " 0.01905834674835205,\n",
       " -0.006711882073432207,\n",
       " -0.007124985568225384,\n",
       " 0.056979432702064514,\n",
       " 0.013075708411633968,\n",
       " -0.013979261741042137,\n",
       " -0.11352594941854477,\n",
       " -0.03491680696606636,\n",
       " 0.013942153193056583,\n",
       " 0.043169185519218445,\n",
       " 0.0068232291378080845,\n",
       " -0.0008051415206864476,\n",
       " 0.006580108776688576,\n",
       " -0.025622187182307243,\n",
       " -0.05044765770435333,\n",
       " 0.05245169252157211,\n",
       " -0.0641460195183754,\n",
       " -0.055498600006103516,\n",
       " -0.007446881849318743,\n",
       " -5.276017327560112e-05,\n",
       " -0.029843736439943314,\n",
       " 0.02432435005903244,\n",
       " 0.06417108327150345,\n",
       " -0.06314094364643097,\n",
       " 0.04182494059205055,\n",
       " 0.042477674782276154,\n",
       " -0.012201766483485699,\n",
       " 0.09175929427146912,\n",
       " -0.021128077059984207,\n",
       " -0.03640036657452583,\n",
       " 0.08847282081842422,\n",
       " 0.11529214680194855,\n",
       " -0.026308033615350723,\n",
       " -0.046475768089294434,\n",
       " 0.04068070650100708,\n",
       " 0.01918831467628479,\n",
       " -0.007530891802161932,\n",
       " -0.07138030230998993,\n",
       " -0.040803324431180954,\n",
       " 0.032108619809150696,\n",
       " -0.0330362431704998,\n",
       " -0.01749657653272152,\n",
       " 0.034432828426361084,\n",
       " 0.048382822424173355,\n",
       " -0.011264889501035213,\n",
       " -0.029528144747018814,\n",
       " 0.006507426034659147,\n",
       " -0.08372807502746582,\n",
       " 0.10181766003370285,\n",
       " -0.041403524577617645,\n",
       " -0.0535278283059597,\n",
       " -0.03146831691265106,\n",
       " 0.006305839400738478,\n",
       " 0.020719777792692184,\n",
       " 0.03841579705476761,\n",
       " 0.0018788420129567385]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Call `get_text_embedding_batch` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = embed_model.get_text_embedding_batch(\n",
    "    [\n",
    "        \"An Amazon SageMaker endpoint is a fully managed resource that enables the deployment of machine learning models\",\n",
    "        \"Sagemaker is integrated with llamaIndex\",\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(embeddings)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
